# CHECKPOINT_PATH=checkpoints/gpt2_345m
# VOCAB_FILE=gpt2-vocab.json
# MERGE_FILE=gpt2-merges.txt
# DATA_PATH=my-gpt2_text_document

GPT_ARGS="--num-layers 1 \
          --hidden-size 10240 \
          --num-attention-heads 16 \
          --seq-length 1024 \
          --max-position-embeddings 1024 \
          --micro-batch-size 4 \
          --global-batch-size 32 \
          --lr 0.00015 \
          --train-iters 500000 \
          --lr-decay-iters 320000 \
          --lr-decay-style cosine \
          --lr-warmup-fraction .01 \
          --dummy-vocab-size 128000 \
          --fp16"

OUTPUT_ARGS="--log-interval 10 \
             --save-interval 500 \
             --eval-interval 100 \
             --eval-iters 10 \
             --checkpoint-activations"

GPUS_PER_NODE=4
# Change for multinode config
MASTER_ADDR=127.0.0.2
MASTER_PORT=6001
NNODES=1
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))

DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT"

PARALLEL_ARGS="--tensor-model-parallel-size 2 --pipeline-model-parallel-size 2"
python -m torch.distributed.launch $DISTRIBUTED_ARGS \
    train_dummy_lm.py \
    $GPT_ARGS \
    $OUTPUT_ARGS \
    $PARALLEL_ARGS \
    --no-bias-dropout-fusion \
    --no-bias-gelu-fusion \
    --no-masked-softmax-fusion